# Copyright (c) 2023 Presto Labs Pte. Ltd.
# Author: leon

import logging
import time
import threading
from concurrent.futures import ThreadPoolExecutor

import requests
import tornado.gen
from google.protobuf import json_format

import coin.exchange.base.kr_rest.private_client_base as privcb
from coin.exchange.uniswap_v2.kr_rest.product import UniswapProduct
from coin.proto.coin_query_pb2 import (
    AccountBalance,
    CurrencyBalance,
)
from coin.exchange.base.kr_rest.rest_client_base import RestUpdate
from xunkemgmt_client.client.api_client import XunkemgmtClient
from coin2.service.web3.order_service_pb2 import (
    QueryAccountBalanceRequestProto, QueryAccountBalanceResponseProto,
    QueryProvideLiquidityRequestProto, QueryProvideLiquidityResponseProto,
    QuerySubmitOrderRequestProto, QuerySubmitOrderResponseProto,
    QueryWithdrawLiquidityRequestProto, QueryWithdrawLiquidityResponseProto
)


class UniswapPrivateParser(object):
    @staticmethod
    def parse_balance(update_msg):
        account_balances = json_format.MessageToDict(update_msg, preserving_proto_field_name=True)
        if account_balances['status'] == 'success':
            currency_balance_list = [
                CurrencyBalance(
                    currency=currency,
                    currency_native=currency,
                    available=balance["available"],
                    hold=balance["hold"],
                    total=balance["hold"] + balance["available"],
                ) for currency, balance in account_balances['currency_to_balance'].items()
            ]
        else:
            currency_balance_list = []
        return AccountBalance(exchange="Uniswap", market_type="Spot", each_balance=currency_balance_list)


class UniswapPrivateClient(privcb.PrivateClientBase):
    ProductType = UniswapProduct

    def __init__(self, wallet_address, og_info):
        privcb.PrivateClientBase.__init__(self)
        self._wallet_address = wallet_address
        self._currency_list = [curr.native_currency for curr in og_info.currencies]
        self._symbol_list = [prod.symbol for prod in og_info.products]
        self._web3_api_client = XunkemgmtClient()
        self._executor = ThreadPoolExecutor(max_workers=2)
        self._fire_lock_map = {symbol: threading.Lock() for symbol in self._symbol_list}

    def _async_account_balance_impl(self):
        if self._executor:
            balance_request = QueryAccountBalanceRequestProto(
                blockchain_name='ethereum',
                mea='Spot.Uniswap.v2',
                wallet_address=self._wallet_address,
                currencies=self._currency_list,
            )
            fut = self._executor.submit(self._web3_api_client.query_account_balance, balance_request)
            return fut
        return None

    @tornado.gen.coroutine
    def async_query_account_balance_impl(self):
        req_timestamp = time.time()
        fut = self._async_account_balance_impl()
        account_balance_response_proto = yield fut
        resp_timestamp = time.time()
        update = RestUpdate(
            req_timestamp,
            resp_timestamp,
            account_balance_response_proto,
            requests.codes.ok,
        )
        update.msg = UniswapPrivateParser.parse_balance(update.msg)
        return update

    def query_fills_impl(self, product):
        raise NotImplementedError()

    def query_list_orders_impl(self, product):
        raise NotImplementedError()

    def _query_submit_impl(self, symbol, side, price, qty):
        lock = self._fire_lock_map[symbol]
        if lock.acquire(False):
            order_request = QuerySubmitOrderRequestProto(
                blockchain_name='ethereum',
                mea='Spot.Uniswap.v2',
                wallet_address=self._wallet_address,
                symbol=symbol,
                price=float(price),
                qty=float(qty),
                side=side.name,
            )
            order_response_proto = self._web3_api_client.query_submit_order(order_request)
            lock.release()
            return json_format.MessageToDict(order_response_proto, preserving_proto_field_name=True)
        else:
            logging.warning(f"order rejected: {symbol}")
            return None

    def query_submit_impl(self, product, side, price, qty):
        symbol = product.native_symbol
        if self._executor:
            return self._executor.submit(self._query_submit_impl, symbol, side, price, qty)
        else:
            return None

    def query_cancel_product_impl(self, product):
        raise NotImplementedError()

    def query_account_balance_impl(self):
        raise NotImplementedError()
